# -*- coding: utf-8 -*-
import sys
import os
from PySide import QtGui
from PySide import QtCore


class Capture(QtGui.QWidget):
    def __init__(self, parent=None):
        super(Capture, self).__init__(parent)
        self.desktop = QtGui.QApplication.desktop()
        self.rubber_band = QtGui.QRubberBand(QtGui.QRubberBand.Rectangle, self)
        self.origin = None
        self.desktop_bg = None
        self.init_ui()

    def init_ui(self):
        width, height = self.desktop.width(), self.desktop.height()
        self.desktop_bg = QtGui.QPixmap.grabWindow(self.desktop.winId(), 0, 0, width, height)
        palette = QtGui.QPalette()
        palette.setBrush(QtGui.QPalette.Background, QtGui.QBrush(self.desktop_bg))
        self.setPalette(palette)
        self.setWindowFlags(QtCore.Qt.CustomizeWindowHint | QtCore.Qt.WindowStaysOnTopHint | QtCore.Qt.FramelessWindowHint)
        self.setCursor(QtCore.Qt.CrossCursor)
        self.resize(width, height)

    def keyPressEvent(self, event):
        print(event)
        if event.key() == QtCore.Qt.Key_Escape:
            self.close()

    def mousePressEvent(self, event):
        self.origin = event.pos()
        self.rubber_band.setGeometry(QtCore.QRect(self.origin, QtCore.QSize()))
        self.rubber_band.show()

    def mouseMoveEvent(self, event):
        self.rubber_band.setGeometry(QtCore.QRect(self.origin, event.pos()).normalized())

    def mouseReleaseEvent(self, event):
        screen_shot = self.desktop_bg.copy(QtCore.QRect(self.rubber_band.pos(), self.rubber_band.size()))
        mouse_button = event.button()
        if mouse_button == QtCore.Qt.LeftButton:
            clipboard = QtGui.QApplication.clipboard()
            clipboard.setPixmap(screen_shot)
        elif mouse_button == QtCore.Qt.RightButton:
            file_name, _ = QtGui.QFileDialog.getSaveFileName(self, '保存', '.', 'Image Files(*.png)')
            if file_name:
                screen_shot.save(file_name, os.path.splitext(file_name)[1])
        self.rubber_band.hide()
        self.close()


if __name__ == '__main__':
    app = QtGui.QApplication(sys.argv)
    capture = Capture()
    capture.show()
    app.exec_()
